from machine import Pin, PWM, UART, FPIOA
from libs.PipeLine import PipeLine, ScopedTiming
from libs.YOLO import YOLOv8
import os, sys, gc, time
import random

# 系统配置参数
CONFIG = {
    "motor1_rpm": 20,         # 电机1基础转速
    "motor2_rpm": 50,         # 电机2基础转速
    "detect_duration": 5,     # 目标检测持续时间（秒）
    "detect_interval": 0.2,   # 检测间隔（秒）
    "stm32_delay": 10,         # STM32动作等待时间（秒）
    "debounce_delay": 0,     # 传感器消抖时间（毫秒）
    "confidence_thresh": 0.3, # 检测置信度阈值
    "uart_timeout": 100       # 串口接收超时（毫秒）
}

# 全局状态变量
system_paused = False  # 系统暂停标志（True=暂停，False=运行）
u1 = None  # UART1对象
u2 = None  # UART2对象
def hardware_init():
    """硬件初始化函数"""
    global fpioa, motor1, motor2, sensor, u1, u2
    # 初始化FPIOA
    fpioa = FPIOA()

    # 配置PWM引脚（电机控制）
    fpioa.set_function(46, FPIOA.PWM2)  # 电机1 PWM引脚
    fpioa.set_function(47, FPIOA.PWM3)  # 电机2 PWM引脚

    # 初始化PWM通道
    motor1 = PWM(2, 1000, 50, enable=True)  # 通道2，1kHz，50%占空比，立即启用
    motor2 = PWM(3, 1000, 50, enable=True)  # 通道3，参数同上

    # 配置遮挡传感器（GPIO6，上拉输入模式）
    fpioa.set_function(6, FPIOA.GPIO6)
    sensor = Pin(6, Pin.IN, Pin.PULL_UP)

    # 初始化串口UART1
    fpioa.set_function(3, FPIOA.UART1_TXD)
    fpioa.set_function(4, FPIOA.UART1_RXD)
    fpioa.set_function(11, FPIOA.UART2_TXD)
    fpioa.set_function(12, FPIOA.UART2_RXD)
    u1 = UART(UART.UART1,9600, timeout=CONFIG["uart_timeout"])
    u2 = UART(UART.UART2,9600)

def ai_init():
    """AI模型初始化函数"""
    global pl, yolo

    # 显示管道配置
    display_mode = "hdmi"               # 显示输出模式
    rgb888p_size = [1280, 720]          # 摄像头采集分辨率
    display_size = [1920, 1080] if display_mode == "hdmi" else [800, 480]  # 显示分辨率

    # 创建图像处理管道
    pl = PipeLine(
        rgb888p_size=rgb888p_size,
        display_size=display_size,
        display_mode=display_mode
    )
    pl.create()  # 初始化管道

    # 初始化YOLOv8目标检测模型
    yolo = YOLOv8(
        task_type="detect",              # 检测任务类型
        mode="video",                    # 视频模式
        kmodel_path="/data/best.kmodel",# 模型文件路径
        labels=["recycle","kitchen","harmful","others"], # 类别标签
        rgb888p_size=rgb888p_size,       # 输入图像尺寸
        model_input_size=[320, 320],     # 模型输入尺寸
        display_size=display_size,       # 显示尺寸
        conf_thresh=CONFIG["confidence_thresh"],  # 置信度阈值
        nms_thresh=0.5,                  # NMS阈值
        max_boxes_num=50                 # 最大检测框数量
        #debug_mode=0                     # 关闭调试模式提升性能
    )
    yolo.config_preprocess()  # 配置预处理参数

def set_motor_speed(motor, rpm):
    """设置电机转速
    Args:
        motor: PWM对象
        rpm: 目标转速（转/分钟）
    """
    pulses_per_rev = 200 * 16  # 200步/圈 * 16细分
    freq = (rpm * pulses_per_rev) / 60  # 计算PWM频率
    motor.freq(int(freq))      # 更新PWM频率

def process_uart_command():
    global system_paused
    # 读取1字节
    data = u1.read()  # 使用read方法读取单个字节

    # 字节数据直接比较
    if data == b'E':
        system_paused = False
        motor1.enable(True)
        motor2.enable(True)
        print("1")
    elif data == b'F':
        system_paused = True
        motor1.enable(False)
        motor2.enable(False)
        print("2")

def sensor_triggered():
    """传感器触发检测（带消抖）"""
    # 第一次检测到低电平
    if sensor.value() == 0:
        start = time.ticks_ms()
        print("YES")
        # 持续检测消抖时间
        while time.ticks_diff(time.ticks_ms(), start) < CONFIG["debounce_delay"]:
            if sensor.value() != 0:
                #print(f"YES")
                return False
        time.sleep(0.5)
        return True
    return False

def detection_phase():
    """执行目标检测流程（持续检测直到获得有效结果）"""
    motor1.enable(False)
    motor2.enable(False)

    total_counter = {}
    label_map = {
        "recycle": b'A',
        "kitchen": b'C',
        "harmful": b'B',
        "others": b'D'
    }

    cycle_start = time.time()
    while (time.time() - cycle_start) < CONFIG["detect_duration"]:
        img = pl.get_frame()

        # 获取检测结果（此时detections是numpy数组）
        detections = yolo.run(img)

        # 绘制检测结果到屏幕
        with ScopedTiming("total", 1):
            yolo.draw_result(detections, pl.osd_img)
            pl.show_image()

        # 解析检测结果的正确方式
        if detections is not None and len(detections) > 0:
            # 假设检测结果结构为：[x1, y1, x2, y2, conf, cls]
            for det in detections:
                try:
                    # 安全获取类别索引
                    if len(det) < 6: continue
                    cls_idx = int(det[5])  # 第6个元素是类别索引
                    # 校验索引有效性
                    if 0 <= cls_idx < len(yolo.labels):
                        label = yolo.labels[cls_idx]
                        if label in label_map:
                            total_counter[label] = total_counter.get(label, 0) + 1

                except Exception as e:
                    print(f"检测数据处理异常: {str(e)}")

        time.sleep(CONFIG["detect_interval"])
        gc.collect()
    if total_counter:
        max_label = max(total_counter, key=total_counter.get)
    else:
        # 随机选择逻辑
        max_label = random.choice(list(label_map.keys()))
        print(f"[WARN] 未检测到有效目标，随机选择: {max_label}")

        # 发送结果（保持原有校验）
    if max_label not in label_map:
        raise RuntimeError(f"最终结果错误: {max_label}")

    final_result = label_map[max_label]
    u1.write(final_result)
    print(f"发送分类指令: {final_result}")

    time.sleep(CONFIG["stm32_delay"])
    u2.write(final_result)

    # 恢复电机
    motor1.enable(True)
    motor2.enable(True)

def main():
    """主程序"""
    # 初始化硬件
    hardware_init()
    set_motor_speed(motor1, CONFIG["motor1_rpm"])
    set_motor_speed(motor2, CONFIG["motor2_rpm"])
    ai_init()

    try:
        while True:
            # 优先处理串口命令
            process_uart_command()

            if not system_paused:
                # 正常操作模式
                if sensor_triggered():
                    print("Obstruction detected!")
                    detection_phase()  # 执行完整检测流程

                    # 等待传感器释放
                    while sensor.value() == 0:
                        time.sleep(0.1)

                    # 恢复电机运转
                    motor1.enable(True)
                    motor2.enable(True)

                    # 防重复触发延时
                    time.sleep(1)
                else:
                    # 维持运转状态
                    time.sleep_ms(50)
            else:
                # 系统暂停状态
                time.sleep(1)

    except KeyboardInterrupt:
        # 安全关闭程序
        motor1.enable(False)
        motor2.enable(False)
        if u1 is not None:
            u1.deinit()  # 释放UART1资源
        if u2 is not None:
            u2.deinit()  # 释放UART2资源
        pl.destroy()
        print("System shutdown")

if __name__ == "__main__":
    main()
